/***********************************************************************/
/* Name:       Xuqing Wu                                               */
/* Email:      xwu@cs.uh.edu                                           */
/* QIL Lab, CS, U. of Houston                                          */
/*                                                                     */
/* All Rights Reserved                                                  */
/* THIS SOURCE CODE IS PROVIDED 'AS IS', WITH NO WARRANTIES WHATSOEVER */
/***********************************************************************/

#include "hdr.h"

#define assigndata2node(plisthead,indx,data) plisthead[idx].data=data

/* example on how to type cast pointer to a function 
typedef double (* deflocalf)(node_t *, int, double *);
typedef double (* defpairf)(node_t *, node_t *, int, int, double *);
*/

node_t *initall(int totalnodes, double **pval, int *pys, int *nbrnum, int **nbridx) 
{
   int      i, j;
   node_t   *p2allnodes, *p2nbr, *ptmp;

   printf("Start init all nodes...");

   p2allnodes = (node_t *) malloc(sizeof(node_t)*totalnodes);
   memset((char *)p2allnodes, 0, sizeof(node_t)*totalnodes);

   for(i=0;i<totalnodes;i++) {
      for(j=0;j<nbrnum[i];j++) {
         p2nbr = &p2allnodes[nbridx[i][j]];
         setneighbor(&p2allnodes[i], p2nbr);
      }
   }

   for(i=0;i<totalnodes;i++) {
      ptmp = &p2allnodes[i];
      if(nbrnum[i] != ptmp->nbrs) {
         printf("ERROR: neighbor data does not match!\n");
      }
      initnode(ptmp,i,pval[i],pys);
   }

   printf("Finish init all nodes\n");

   return p2allnodes;
}

/* find message value from box#boxid saved by psrc for y state
 */
double findmsg(node_t *pnode, node_t *psrc, int y, int boxid)
{
   double      msgval = 1;
   message_t   **mbox;
   message_t   *pmsg;
   nodehash_t  *phash;

   if(boxid == BOXID1) {
      mbox = pnode->mbox1;
   } else {
      mbox = pnode->mbox2;
   }

   HASH_FIND_PTR(pnode->nhash, &psrc, phash);
   if(! phash)
      printf("ERROR: cannot find neighbor in hash table\n");
   pmsg = &mbox[y][phash->nid];
   if(pmsg->source != psrc)
      printf("ERROR: findmsg with inconsistent message source!\n");
   msgval = pmsg->val;

   return msgval;
}

/* product of message transferred to pnode from all neighbors
 * except pnot with state of y from specific boxid
 */
double prodnbrmsgs(node_t *pnode, node_t *pnot, int y, int boxid)
{
   double      msgval = 1.0;
   node_t      *pcurnode;
   nodehash_t  *tmphash;

   for(tmphash=pnode->nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
      pcurnode = tmphash->key;
      if(pcurnode != pnot) {
         msgval = msgval * pow(findmsg(pnode, pcurnode, y, boxid),NE);
      }
   }

   return msgval;
}

/* find the normlization factor K for Mji(x_i)
 */
double normlizerK(node_t *pnode, node_t *psrc, int boxid)
{
   int      i;
   double   k, total = 0;

   for(i=0;i<NUMOFSTATE;i++) {
      total += findmsg(pnode, psrc, i, boxid);
   }
   
   k = 1 / total;
   return k;
}

/* send message from one node psrc to another node ptarget with state yi
 */
void sendmsg2(node_t *ptarget, node_t *psrc, int boxid, int yi, double msgval,
      int idx)
{
   message_t **mbox, *pmsg;
   nodehash_t *phash;
   int i;

   if(boxid == BOXID1) {
      mbox = ptarget->mbox1;
   } else {
      mbox = ptarget->mbox2;
   }
   
   HASH_FIND_PTR(ptarget->nhash, &psrc, phash);
   pmsg = &mbox[yi][phash->nid];
   if(pmsg->source != psrc)
      printf("ERROR: sendmsg2 hash table is inconsistent!\n");
   pmsg->val = msgval;
   pmsg->idx = idx;
}

/* transfer message from pnode to a neighbor for a specific state i
 */
void mp2node4i(node_t *psrc, node_t *ptarget, int yi, double *para1, 
      double *para2, int boxid)
{
   int         i;
   double      prodmsg=0.0, invmsg=0.0;
   double      e1=0.0, e2=0.0, ep=0.0, etotal=0.0;

   for(i=0;i<NUMOFSTATE;i++) {
      e1 = psrc->clocal(psrc, i, para1);
      e2 = psrc->cpair(psrc, ptarget, i, yi, para2);
      ep = exp(-1/NE*e2-e1);
      prodmsg = prodnbrmsgs(psrc, ptarget, i, boxid);
      invmsg = findmsg(psrc,ptarget,i,boxid);
      etotal += ep*prodmsg/(pow(invmsg,1-NE)+EPSILON);
   }
   etotal = etotal*normlizerK(ptarget, psrc, boxid);
   
   if(boxid == BOXID1) {
      //sendmsg2(ptarget, psrc, BOXID1, yi, etotal, -1);
      sendmsg2(ptarget, psrc, BOXID2, yi, etotal, -1);
   } else {
      //sendmsg2(ptarget, psrc, BOXID2, yi, etotal, -1);
      sendmsg2(ptarget, psrc, BOXID1, yi, etotal, -1);
   }
}

/* transfer message from pnode to a neighbor for a specific state i in 
 * max-product mode
 */
void mpmax2node4i(node_t *psrc, node_t *ptarget, int yi, double *para1, 
      double *para2, int boxid)
{
   int         i, idx;
   double      prodmsg=0.0, invmsg=0.0;
   double      e1=0.0, e2=0.0, ep=0.0, etotal=0.0, emax=0.0;

   for(i=0;i<NUMOFSTATE;i++) {
      e1 = psrc->clocal(psrc, i, para1);
      e2 = psrc->cpair(psrc, ptarget, i, yi, para2);
      ep = exp(-1/NE*e2-e1);
      prodmsg = prodnbrmsgs(psrc, ptarget, i, boxid);
      invmsg = findmsg(psrc,ptarget,i,boxid);
      etotal = ep*prodmsg/(pow(invmsg,1-NE)+EPSILON);
      if(etotal > emax) {
         emax = etotal;
         idx = i;
      }
   }
   emax = emax*normlizerK(ptarget, psrc, boxid);
   
   if(boxid == BOXID1) {
      //sendmsg2(ptarget, psrc, BOXID1, yi, emax, idx);
      sendmsg2(ptarget, psrc, BOXID2, yi, emax, idx);
   } else {
      //sendmsg2(ptarget, psrc, BOXID2, yi, emax, idx);
      sendmsg2(ptarget, psrc, BOXID1, yi, emax, idx);
   }
}


/* transfer message from every node to all of its neighbors for 
 * all states of y
 */
void runbpfd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   node_t      *pnode;
   nodehash_t  *tmphash;

   for(i=0;i<totalnodes;i++) {
      for(tmphash=phead[i].nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnode = tmphash->key;
         for(j=0;j<NUMOFSTATE;j++) {
            mp2node4i(&phead[i], pnode, j, para1, para2, boxid);
         }
      }
   }
}

void runbpbd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   node_t      *pnode;
   nodehash_t  *tmphash;

   for(i=totalnodes-1;i>=0;i--) {
      for(tmphash=phead[i].nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnode = tmphash->key;
         for(j=0;j<NUMOFSTATE;j++) {
            mp2node4i(&phead[i], pnode, j, para1, para2, boxid);
         }
      }
   }
}

/* transfer message from every node to all of its neighbors for 
 * all states of y in max product mode
 */
void runmaxbpfd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodehash_t  *tmphash;
   node_t      *pnode;

   for(i=0;i<totalnodes;i++) {
      //printf("Forward max-msg for node#%d\n", i); 
      for(tmphash=phead[i].nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnode = tmphash->key;
         for(j=0;j<NUMOFSTATE;j++) {
            mpmax2node4i(&phead[i], pnode, j, para1, para2, boxid);
         }
      }
   }
}

void runmaxbpbd(node_t *phead, int totalnodes, double *para1, double *para2,
      int boxid)
{
   int         i, j;
   nodehash_t  *tmphash;
   node_t      *pnode;

   for(i=totalnodes-1;i>=0;i--) {
      //printf("Backward max-msg for node#%d\n", i);
      for(tmphash=phead[i].nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnode = tmphash->key;
         for(j=0;j<NUMOFSTATE;j++) {
            mpmax2node4i(&phead[i], pnode, j, para1, para2, boxid);
         }
      }
   }
}

/* compute belief value bi for a node
 */
void computebi(node_t *pnode, double *para1, double *pk, int boxid)
{
   int         i, j;
   double      total=0.0, ei=0.0, prodmsg=1.0;
   message_t   **mbox;

   if(boxid == BOXID1) {
      mbox = pnode->mbox1;
   } else {
      mbox = pnode->mbox2;
   }
   
   for(i=0;i<NUMOFSTATE;i++) {
      prodmsg = 1.0;
      ei = exp(-pnode->clocal(pnode, i, para1));
      for(j=0;j<pnode->nbrs;j++) {
         prodmsg = prodmsg*pow(mbox[i][j].val, NE);
      }
      pnode->bi[i] = ei*prodmsg;
      total += pnode->bi[i];
   }
   for(i=0;i<NUMOFSTATE;i++) {
      *pk = 1 / total;
      pnode->bi[i] = pnode->bi[i] / total;
   }
}

/* compute pair belief bij for node i
 */

void computebij(node_t *pnode, double *para1, double *para2, double pk, 
      int boxid)
{
   int         i, j, t, idx=0;
   double      eij=0.0, ei=0.0, ej=0.0, prodmsg1=1.0, prodmsg2=1.0;
   double      invmsg1 = 1.0, invmsg2 = 1.0;
   nodehash_t  *tmphash;
   node_t      *pnbr;
   message_t   **mboxi, **mboxj;

   for(i=0;i<NUMOFSTATE;i++) {
      for(j=0;j<NUMOFSTATE;j++) {
         idx = 0;
         for(tmphash=pnode->nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
            pnbr = tmphash->key;
            if(boxid == BOXID1) {
               mboxi = pnode->mbox1;
               mboxj = pnbr->mbox1;
            } else {
               mboxi = pnode->mbox2;
               mboxj = pnbr->mbox2;
            }

            ei = pnode->clocal(pnode, i, para1);
            ej = pnbr->clocal(pnbr, j, para1);
            eij = pnode->cpair(pnode, pnbr, i, j, para2);
            for(t=0;t<pnode->nbrs;t++) {
               if(mboxi[i][t].source != pnbr) {
                  prodmsg1 *= pow(mboxi[i][t].val, NE);
               } else {
                  invmsg1 = pow(mboxi[i][t].val, 1-NE);
               }
            }
            for(t=0;t<pnbr->nbrs;t++) {
               if(mboxj[j][t].source != pnode) {
                  prodmsg2 *= pow(mboxj[j][t].val, NE);
               } else {
                  invmsg2 = pow(mboxj[j][t].val, 1-NE);
               }
            }
   
            pnode->bij[j+NUMOFSTATE*i][idx].val = pk*exp(-1/NE*eij-ei-ej) *
               prodmsg1/(invmsg1+EPSILON)*prodmsg2/(invmsg2+EPSILON);
            pnode->bij[j+NUMOFSTATE*i][idx].pnbr = pnbr;

            idx++;
         }
      }
   }
}

/* forward back tracing for the max-product result
 */
void fdbt4bp(node_t *phead, int totalnodes, int boxid)
{
   int         i, j, status, maxidx=0;
   double      maxbi=0.0;
   node_t      *pnode, *pnbr;
   nodehash_t  *tmphash;
   message_t   **mbox;

   printf("Start forward backtracing for blief value\n");

   pnode = &phead[0];
   for(i=0;i<NUMOFSTATE;i++) {
      if(pnode->bi[i] > maxbi) {
         maxbi = pnode->bi[i];
         maxidx = i;
      }
   }
   pnode->ymax = maxidx;

   for(i=0;i<totalnodes;i++) {
      pnode = &phead[i];
      if(pnode->ymax < 0) {
         printf("ERROR: the current node is not assigned in backtracing!\n");
      }  
      maxidx = pnode->ymax;

      if(boxid == BOXID1) {
         mbox = pnode->mbox1;
      } else {
         mbox = pnode->mbox2;
      }
  
      //printf("Search forward max index for node#%d\n",i);
      for(tmphash=pnode->nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnbr = tmphash->key;
         status = 0;
         for(j=0;j<pnode->nbrs;j++) {
            if(mbox[maxidx][j].source == pnbr) {
               status = 1;
               if(pnbr->ymax < 0) {
                  pnbr->ymax = mbox[maxidx][j].idx;
               }
               break;
            }
         }
         if(! status) {
            printf("ERROR: backtracing cannot find idx for %p\n", pnbr);
         }
      }
   }

   printf("Finish forward backtracing for belief value\n");
}

/* backward back tracing for the max-product result
 */
void bdbt4bp(node_t *phead, int totalnodes, int boxid)
{
   int         i, j, status, maxidx=0;
   double      maxbi=0.0;
   node_t      *pnode, *pnbr;
   nodehash_t  *tmphash;
   message_t   **mbox;

   printf("Start backward backtracing for blief value\n");
   pnode = &phead[totalnodes-1];
   for(i=0;i<NUMOFSTATE;i++) {
      if(pnode->bi[i] > maxbi) {
         maxbi = pnode->bi[i];
         maxidx = i;
      }
   }
   pnode->ymax = maxidx;

   for(i=totalnodes-1;i>=0;i--) {
      pnode = &phead[i];
      if(pnode->ymax < 0) {
         printf("ERROR: the current node is not assigned in backtracing!\n");
      }  
      maxidx = pnode->ymax;

      if(boxid == BOXID1) {
         mbox = pnode->mbox1;
      } else {
         mbox = pnode->mbox2;
      }
  
      //printf("Search backward max index for node#%d\n",i);
      for(tmphash=pnode->nhash;tmphash!=NULL;tmphash=tmphash->hh.next) {
         pnbr = tmphash->key;
         status = 0;
         for(j=0;j<pnode->nbrs;j++) {
            if(mbox[maxidx][j].source == pnbr) {
               status = 1;
               if(pnbr->ymax < 0) {
                  pnbr->ymax = mbox[maxidx][j].idx;
               }
               break;
            }
         }
         if(! status) {
            printf("ERROR: backtracing cannot find idx for %p\n", pnbr);
         }
      }
   }

   printf("Finish backward backtracing for blief value\n");
}

void sortbybi(node_t *phead, int totalnodes)
{
   int      i, j;
   double   yval;
   node_t   *pnode;

   for(i=0;i<totalnodes;i++) {
      yval = 0.0;
      pnode = &phead[i];
      for(j=0;j<NUMOFSTATE;j++) {
         if(pnode->bi[j] > yval) {
            pnode->ymax = j;
            yval = pnode->bi[j];
         }
      }
   }
}

void freeallnodes(node_t *phead, int totalnodes)
{
   int i, j;
   node_t *pnode;
   nodehash_t *phash;

   printf("Start freeing mem for all nodes...");
   for(i=0;i<totalnodes;i++) {
      pnode = &phead[i];

      /* free message box from each node
       */
      for(j=0;j<NUMOFSTATE;j++) {
         free(pnode->mbox1[j]);
         free(pnode->mbox2[j]);
      }

      /* free bij from each node
       */
      for(j=0;j<NUMOFSTATE*NUMOFSTATE;j++) {
         free(pnode->bij[j]);
      }

      /* free hash table for each node
       */
      while(pnode->nhash) {
         phash = pnode->nhash;
         HASH_DEL(pnode->nhash, phash);
         free(phash);
      }
   }

   free(phead);
   printf("Finish freeing mem for all nodes\n");
}

